# encoding=utf-8
import os
import json
import tqdm
from collections import OrderedDict
import config
from util.get_model_instance import get_model_instance

if __name__ == '__main__':
    model_name = config.MODEL_NAME
    # batch_dir：不同模型生成数据的目录
    batch_dir = '../mydata/'+model_name+'_generations_'+config.BATCH_DIR_SUFFIX+'/'
    # 本步骤输入输出文件路径
    input_file = 'machine_generated_instructions.jsonl'
    output_file = 'machine_generated_instances.jsonl'

    # 需要让LLM遵循的指令
    tasks = []
    with open(os.path.join(batch_dir, input_file)) as fin:
        lines = fin.readlines()
        for line in lines:
            data = json.loads(line)
            tasks.append(data)

    output_path = os.path.join(batch_dir, output_file)
    existing_requests = {}

    # 加载已经得到的输出
    if os.path.exists(output_path):
        with open(output_path, encoding='utf-8') as fin:
            for line in tqdm.tqdm(fin):
                try:
                    data = json.loads(line)
                    existing_requests[data["instruction"]] = data   # 这是个字典，key是instruction，value是整行data
                except:
                    pass
        print(f"Loaded {len(existing_requests)} existing requests")

    model_instance = get_model_instance(model_name)
    progress_bar = tqdm.tqdm(total=len(tasks))
    with open(output_path, "w", encoding='utf-8') as fout:
        for task in tasks:
            if task["instruction"] in existing_requests:
                data = existing_requests[task["instruction"]]
                data = OrderedDict(
                    (k, data[k]) for k in \
                    ["instruction", "output",
                     "most_similar", "avg_similarity_score"]
                )
                fout.write(json.dumps(data, ensure_ascii=False) + "\n")
            else:
                data = task
                prompt = config.HANDWRITE_PROMPT_OUTPUT + '\n' + '请根据以上背景知识，用中文回答：' + task["instruction"] + '注意不用特别指出基于背景知识，并且注意不要用英文回答问题。'
                data["output"] = model_instance.get_raw_output(prompt)

                data = OrderedDict(
                    (k, data[k]) for k in \
                    ["instruction", "output",
                     "most_similar", "avg_similarity_score"]
                )

                if progress_bar.n % 100 == 0:
                    print(data)

                fout.write(json.dumps(data, ensure_ascii=False) + "\n")
            progress_bar.update(1)
